package com.yth.utils;


import lombok.extern.slf4j.Slf4j;

import java.lang.ref.ReferenceQueue;
import java.lang.ref.WeakReference;
import java.security.AccessController;
import java.security.PrivilegedAction;
import java.util.Objects;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.ThreadFactory;
import java.util.concurrent.atomic.AtomicBoolean;


/**
 * Allows a way to register some {@link Runnable} that will executed once there are no references to an {@link Object}
 * anymore.
 * <p>
 * 改自netty的ObjectCleaner。
 *
 * @author yutianhong
 * @version 1.0
 * @since 2025/3/6 17:28
 */
@Slf4j
public final class ObjectCleaner {
    private static final int REFERENCE_QUEUE_POLL_TIMEOUT_MS = 10000;

    // This will hold a reference to the AutomaticCleanerReference which will be removed once we called cleanup()
    private static final Set<AutomaticCleanerReference> LIVE_SET = ConcurrentHashMap.newKeySet();
    private static final ReferenceQueue<Object> REFERENCE_QUEUE = new ReferenceQueue<>();
    private static final AtomicBoolean CLEANER_RUNNING = new AtomicBoolean(false);

    private static final ExecutorService CLEANER_EXECUTOR = createCleanerExecutor();
    private static final Runnable CLEANER_TASK = new Runnable() {
        @Override
        public void run() {
            boolean interrupted = false;
            //noinspection ConditionalBreakInInfiniteLoop
            while (true) {
                // Keep on processing as long as the LIVE_SET is not empty and once it becomes empty
                // See if we can let this thread complete.
                while (!LIVE_SET.isEmpty()) {
                    final ObjectCleaner.AutomaticCleanerReference reference;
                    try {
                        reference = (AutomaticCleanerReference) REFERENCE_QUEUE.remove(REFERENCE_QUEUE_POLL_TIMEOUT_MS);
                    } catch (InterruptedException ex) {
                        // Just consume and move on
                        interrupted = true;
                        continue;
                    }
                    if (reference != null) {
                        try {
                            reference.cleanup();
                        } catch (Throwable ignored) {
                            // ignore exceptions, and don't log in case the logger throws an exception, blocks, or has
                            // other unexpected side effects.
                        }
                        LIVE_SET.remove(reference);
                    }
                }
                CLEANER_RUNNING.set(false);
                // Its important to first access the LIVE_SET and then CLEANER_RUNNING to ensure correct
                // behavior in multithreaded environments.
                //noinspection ConstantValue
                if (LIVE_SET.isEmpty() || !CLEANER_RUNNING.compareAndSet(false, true)) {
                    // There was nothing added after we set STARTED to false or some other cleanup Thread
                    // was started already so its safe to let this Thread complete now.
                    break;
                }
            }
            if (interrupted) {
                // As we caught the InterruptedException above we should mark the Thread as interrupted.
                Thread.currentThread().interrupt();
            }
        }
    };

    private static ExecutorService createCleanerExecutor() {
        ThreadFactory threadFactory = new ThreadFactory() {
            private final ThreadFactory defaultFactory = Executors.defaultThreadFactory();

            @Override
            public Thread newThread(Runnable r) {
                Thread cleanupThread = defaultFactory.newThread(r);
                cleanupThread.setName(ObjectCleaner.class.getSimpleName() + "Thread");
                cleanupThread.setPriority(Thread.MIN_PRIORITY);
                // Set to null to ensure we not create classloader leaks by holding a strong reference to the inherited
                // classloader.
                // See:
                // - https://github.com/netty/netty/issues/7290
                // - https://bugs.openjdk.java.net/browse/JDK-7008595
                AccessController.doPrivileged((PrivilegedAction<Void>) () -> {
                    cleanupThread.setContextClassLoader(null);
                    return null;
                });
                // Mark this as a daemon thread to ensure that we the JVM can exit if this is the only thread that is
                // running.
                cleanupThread.setDaemon(true);
                return cleanupThread;
            }
        };
        ExecutorService cleanerExecutor = Executors.newSingleThreadExecutor(threadFactory);
        ThreadUtils.addShutdownHook(() -> ThreadUtils.shutdownThreadPool(cleanerExecutor, log));
        return cleanerExecutor;
    }

    /**
     * Register the given {@link Object} for which the {@link Runnable} will be executed once there are no references
     * to the object anymore.
     * <p>
     * This should only be used if there are no other ways to execute some cleanup once the Object is not reachable
     * anymore because it is not a cheap way to handle the cleanup.
     */
    public static void register(Object object, Runnable cleanupTask) {
        AutomaticCleanerReference reference = new AutomaticCleanerReference(object, cleanupTask);
        // Its important to add the reference to the LIVE_SET before we access CLEANER_RUNNING to ensure correct
        // behavior in multithreaded environments.
        LIVE_SET.add(reference);
        if (CLEANER_RUNNING.compareAndSet(false, true)) {
            CLEANER_EXECUTOR.execute(CLEANER_TASK);
        }
    }

    public static int getLiveSetCount() {
        return LIVE_SET.size();
    }

    private ObjectCleaner() {
        // Only contains a static method.
    }

    private static final class AutomaticCleanerReference extends WeakReference<Object> {
        private final Runnable cleanupTask;

        AutomaticCleanerReference(Object referent, Runnable cleanupTask) {
            super(referent, REFERENCE_QUEUE);
            this.cleanupTask = Objects.requireNonNull(cleanupTask);
        }

        void cleanup() {
            cleanupTask.run();
        }

        @Override
        public Thread get() {
            return null;
        }

        @Override
        public void clear() {
            LIVE_SET.remove(this);
            super.clear();
        }
    }
}
